# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
# This file is part of ByteQC.
#
# Licensed under the Apache License, Version 2.0 (the "License")
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https: // www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf import gto
from pyscf import scf
from byteqc import cucc
from pyscf import cc
from time import time
import cupy
import os

device = eval(os.getenv('DEFAULT_DEVICE', '0'))

for basis in ['ccpvdz', 'ccpvtz', 'ccpvqz']:
    isBreak = False
    for nwater in (1, 1, 2, 3, 4, 6, 8, 10, 12, 14, 16,
                   20, 24, 28, 32, 40, 48, 56, 64):
        print('Basis:%s Nwater:%d Device:%s' % (
            basis, nwater, 'GPU' if device == 0 else 'CPU'))
        mol = gto.M()
        mol.atom = '''
            O -6.31504583 -6.96366310 -2.72627950
            H -5.79312801 -6.26504898 -2.37597966
            H -7.14357233 -6.71134567 -2.32921052
            O -8.45501137 -4.28040600 4.57241678
            H -8.41998100 -3.59314179 3.90820169
            H -9.00186634 -3.98211813 5.31658316
            O -3.34254265 -7.53874588 -4.81753922
            H -3.37715197 -8.40490437 -4.44185734
            H -4.16632652 -7.42499590 -5.26726389
            O -6.00126457 2.73719716 4.28764105
            H -5.07266521 2.91264296 4.07020903
            H -6.37168217 2.60768437 3.37017345
            O -3.32503390 4.33411789 -7.89316702
            H -2.59695578 4.66874695 -7.44677067
            H -3.58132029 5.06191063 -8.46547031
            O -10.86732292 2.57142377 6.23482895
            H -10.30360985 1.91717947 6.68361759
            H -10.81357193 3.34719133 6.83979940
            O -4.04160786 -6.00639248 2.82591510
            H -3.95338082 -5.06667852 2.94704604
            H -5.01617479 -6.24317312 2.96309853
            O 1.37160087 -0.38430426 2.60069180
            H 1.23877645 -0.39831364 3.55441070
            H 0.73124522 -0.91313595 2.24173427
            O 0.16538815 3.71141911 1.34374118
            H 0.13583852 3.45060563 2.28717899
            H -0.65650594 3.40389252 1.04012930
            O -4.60142469 -3.04783869 3.58867168
            H -5.00832415 -3.60508657 4.28174829
            H -5.01135921 -2.13793111 3.67012882
            O -8.47128773 5.27224493 -5.48457670
            H -8.37912846 6.16458035 -5.07643795
            H -9.21504116 4.93016529 -5.02637339
            O -2.40646720 1.20477355 -1.77750456
            H -2.49952078 0.55142862 -1.11293638
            H -3.14344025 1.09464312 -2.35185099
            O 2.99062514 0.08082870 -5.12260771
            H 3.43661666 0.89750516 -4.74682856
            H 2.50228548 -0.40059093 -4.42516851
            O -2.17296195 2.63227940 0.51625723
            H -2.32637024 2.28041053 -0.40385184
            H -2.70722580 3.42956281 0.66376978
            O -0.16663843 6.50960827 -3.78184891
            H -0.62454700 7.07884359 -4.35806751
            H -0.81140471 6.08111143 -3.13051558
            O -8.49094486 7.69357872 8.23789406
            H -9.03492355 8.28426933 7.74227190
            H -7.68518734 8.22941208 8.26448250
            O -4.60740566 -2.26502395 -2.95746398
            H -3.72664261 -2.64869761 -2.99651170
            H -5.14192247 -2.80897188 -3.53671670
            O -5.66186476 -9.40780067 -3.66805387
            H -5.46084690 -9.33606815 -4.59570074
            H -5.68517351 -8.49637032 -3.35685420
            O -8.64172840 -6.90924168 -1.00803602
            H -8.22439480 -6.27162790 -0.41556787
            H -9.29443645 -6.46964645 -1.45871425
            O -6.07779980 -4.89796257 5.49510241
            H -6.96759605 -4.64862728 5.15916824
            H -6.02043962 -5.87251663 5.58129311
            O 2.34769464 -3.36739612 -6.01602077
            H 2.04887176 -2.44794869 -6.21551323
            H 1.44258010 -3.85883570 -6.05587196
            O -3.76762342 -7.55055523 0.48845601
            H -3.15312243 -7.25187683 -0.16663116
            H -3.77164173 -6.90572596 1.22618222
            O -3.43009281 1.64490891 4.79866838
            H -4.03584433 1.57588899 5.53983307
            H -3.34224772 2.55351996 4.67301035
            O -0.46168166 -6.12163162 1.09510243
            H 0.10668758 -5.51180315 0.71371299
            H -0.14936532 -7.03945589 0.97468501
            O -6.50911379 -0.79509622 -8.81709671
            H -7.34802055 -1.32536137 -9.07578087
            H -6.32388306 -0.16061966 -9.53934193
            O 7.49520111 -4.91360664 -1.69814694
            H 7.40926695 -4.66315508 -0.78496742
            H 7.69696808 -4.07921362 -2.07061100
            O 3.90259004 -2.11219835 2.64090061
            H 3.17503476 -1.48133922 2.79959226
            H 3.89378977 -2.14871716 1.62508380
            O -0.68540132 9.46042824 2.36137581
            H -0.05739218 8.75185108 2.42537689
            H -1.25110734 9.28205872 1.62822235
            O -0.18799661 7.78617859 -6.63894939
            H 0.24250983 7.51276350 -7.39614630
            H -0.73031974 8.50955963 -6.82543325
            O 4.54612160 7.51634836 0.64674801
            H 5.39381075 8.05878639 0.61367309
            H 4.68234873 7.04517221 1.48108900
            O 3.73654199 0.48572090 -1.15618587
            H 3.20211983 1.06745338 -0.67476559
            H 4.22656679 0.95676386 -1.84796762
            O -2.21734905 8.33113861 0.26419133
            H -3.12351322 8.30922318 0.39998388
            H -1.93985915 7.40440893 0.32687563
            O 4.05333185 -1.98103845 -0.02185570
            H 4.00769091 -1.10421813 -0.39931673
            H 3.65147948 -2.51237321 -0.69508928
            O 1.90192807 -0.73786831 -2.80972743
            H 2.45116353 -0.14531421 -2.29062319
            H 1.17371511 -0.94052315 -2.24266744
            O -4.14224911 0.26880389 -3.47064471
            H -4.90046406 0.75393385 -3.82656622
            H -4.39947176 -0.62374830 -3.20331049
            O -0.72952807 -1.66858363 6.68440771
            H -0.91604143 -1.57507992 7.63071060
            H -1.56143308 -1.84061241 6.26524401
            O -2.27017903 -6.87003994 -2.08701849
            H -2.97441816 -6.76624346 -2.71853900
            H -1.83610511 -7.74420929 -2.27279305
            O 2.75970435 -7.42528820 -8.08861637
            H 3.18323731 -7.34467363 -7.22243023
            H 1.87388873 -7.76346064 -7.93528605
            O -2.25251484 -2.58540106 4.43371773
            H -1.48194635 -2.76315689 3.86267686
            H -3.13204074 -2.77696514 4.09514523
            O 6.79196739 1.08504128 6.35011005
            H 6.49984837 0.16109729 6.31747103
            H 6.47544050 1.62105608 5.58384228
            O 5.49075556 5.64833307 2.51119447
            H 4.73815966 5.30496264 2.96629930
            H 5.66859770 5.10605860 1.74727118
            O 2.94602489 1.66261744 -11.18212891
            H 2.38359404 2.36589193 -11.38939476
            H 2.37403274 0.88985014 -11.07854366
            O -0.73105466 3.30720711 9.47586536
            H -0.46115825 3.61055779 8.62746811
            H -1.22229552 2.49213815 9.33124828
            O -2.29456806 2.04934812 -5.04749346
            H -2.26730871 1.84387732 -5.96944904
            H -2.83287001 1.40745139 -4.59817171
            O 1.78723109 3.34286904 -1.94128072
            H 1.90705061 3.89254904 -1.13334179
            H 0.92618746 2.91742206 -1.82573462
            O 1.19450939 7.50113201 3.06086278
            H 0.77951699 6.69345808 2.85593605
            H 2.01779318 7.18451595 3.45678997
            O 0.25740579 3.77765775 4.20728016
            H -0.25663465 4.22898912 4.86591101
            H 0.80240744 3.18302083 4.75533628
            O 2.08474922 7.29129648 10.03464317
            H 1.32062173 7.13611937 9.52436733
            H 2.82533002 7.41345215 9.43700218
            O 6.13780832 -3.29400706 -4.73900652
            H 6.34290266 -3.81580305 -5.48040247
            H 6.05682564 -2.42694092 -5.18074989
            O 1.31295585 -4.30244303 0.10245909
            H 1.94986438 -4.48227501 0.80672282
            H 1.73553848 -4.51202154 -0.76760751
            O 5.76587772 -10.96664143 1.63199544
            H 4.84435844 -11.00938797 1.62436080
            H 5.89807510 -10.19888783 1.12287343
            O 5.30121088 -1.12653327 6.33600140
            H 4.53303432 -0.59858799 6.61539459
            H 5.25727987 -0.92806637 5.40444708
            O 6.39001274 4.81748056 -6.36224413
            H 6.31594324 4.06403399 -6.94876432
            H 5.47775412 4.88759518 -5.97809219
            O 7.21732378 -0.74047625 0.21252051
            H 6.57920837 -0.37624958 -0.44805548
            H 7.13562298 0.08063880 0.73319989
            O 5.90417480 3.87129521 0.18074951
            H 5.21987009 4.09243488 -0.45769536
            H 6.62958002 4.44024515 -0.09421976
            O -2.98690438 0.55685866 2.40691853
            H -2.52341366 1.29267693 2.00923681
            H -3.19699550 0.77163476 3.30854988
            O -0.31108531 -1.99076760 -1.78423977
            H -0.97196662 -1.71122682 -1.10594165
            H 0.16937001 -2.67924809 -1.37874877
            O 7.31231642 -3.75379443 0.87101680
            H 7.06028128 -2.97903299 0.28558564
            H 7.51060200 -3.42906880 1.75559950
            O 1.21826243 -0.85190046 5.17803812
            H 0.38258818 -1.00604582 5.63397503
            H 1.66010559 -0.32027137 5.83107996
            O -1.06496561 5.10914516 6.14282465
            H -1.87265229 5.07118177 6.69338799
            H -0.69557017 5.98683119 6.06865263
            O -2.22679615 8.84171581 -2.69012761
            H -2.21766090 8.00881004 -2.25659990
            H -1.50517571 9.33919430 -2.35033488
            O 4.01270247 2.12188435 -3.72990346
            H 4.88136101 2.60390139 -3.73689842
            H 3.45489883 2.55304313 -3.05991793
            O 1.69374335 4.52970505 -4.40602541
            H 1.04253161 5.26815844 -4.25164032
            H 1.78998995 4.13418150 -3.51006246
            O 10.43322468 -0.98995972 0.03126749
            H 10.18956375 -0.70881796 0.94120270
            H 9.63492584 -1.43550730 -0.22714230'''
        mol.atom = '\n'.join(mol.atom.split('\n')[1:nwater * 3 + 1])
        mol.basis = basis
        mol.build()
        print("nao:", mol.nao, "nelec:", mol.nelec[0])
        rhf = scf.RHF(mol)
        rhf.conv_tol = 1e-14
        rhf.verbose = 0
        rhf.max_cycle = 0
        rhf.scf()
        rhf = rhf.density_fit()

        def run(mycc):
            mcc = mycc.CCSD(rhf)
            mcc.conv_tol = 1e-14
            mcc.max_cycle = 0
            mcc.ccsd()
            mcc.verbose = 7
            if mycc == cucc:
                from byteqc.cucc.ccsd_t import kernel
                cupy.cuda.Device().synchronize()
            else:
                from pyscf.cc.ccsd_t import kernel
            print("Cal CCSD(T)")
            start = time()
            e = kernel(mcc, mcc.ao2mo())
            if mycc == cucc:
                cupy.cuda.Device().synchronize()
            t = time() - start
            print("CCSD(T) time:", t)
            return e, t

        if device == 0:
            e, t = run(cucc)
            cupy.get_default_memory_pool().free_all_blocks()
            print("GPU:", e, "CCSD(T) time:", t)
        else:
            e, t = run(cc)
            print("CPU:", e, "CCSD(T) time:", t)
        print("\n")
        if t > 10000:
            break
